/*
 * LiME - Linux Memory Extractor
 * Copyright (c) 2011-2012 Digital Forensics Solutions, LLC
 *
 * Contact email:  lime@digitalforensicssolutions.com
 *
 * Author:
 * Joe Sylve       - joe@digitalforensicssolutions.com, @jtsylve
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
 * General Public License for more details. 
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA 
 */

#include "lime.h"

// This file
static int write_lime_header(struct resource *);
static int write_padding(size_t);
static int write_range(struct resource *);
static int write_vaddr(void *, size_t);
static int setup(void);
static void cleanup(void);
static int init(void);

// External
extern int write_vaddr_disk(void *, size_t);
extern int write_count( int, int );
extern int setup_disk(void);
extern void cleanup_disk(void);

static char * format = 0;
static int mode = 0;
static int method = 0;
static int pages_count = 0;
static char zero_page[PAGE_SIZE];

char *   path = 0;
char * rcpath = 0;
int dio = 1;
int port = 0;
int part_no = 0;

extern struct resource iomem_resource;

module_param( path,   charp, S_IRUGO );
module_param( rcpath, charp, S_IRUGO );
module_param( dio,    bool,  S_IRUGO );
module_param( format, charp, S_IRUGO );

int init_module (void)
{
    part_no = 0;
    pages_count = 0;
	if( !path ) {
		DBG("No path parameter specified");
		return -EINVAL;
	}
	if( !rcpath ) {
        DBG("No rc path parameter specificed");
        return -EINVAL;
	}

	if(!format) {
		DBG("No format parameter specified");
		return -EINVAL;
	}
	
	DBG("Parameters");
	DBG("  PATH: %s", path);
    DBG("  RCPATH: %s", rcpath);
	DBG("  DIO: %u", dio);
	DBG("  FORMAT: %s", format);

	memset(zero_page, 0, sizeof(zero_page));

	if (!strcmp(format, "raw")) mode = LIME_MODE_RAW;
	else if (!strcmp(format, "lime")) mode = LIME_MODE_LIME;
	else if (!strcmp(format, "padded")) mode = LIME_MODE_PADDED;
	else {
		DBG("Invalid format parameter specified.");
		return -EINVAL;
	}
    method = LIME_METHOD_DISK;
	return init();
}

static int init() {
	struct resource *p;
	int err = 0;
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2,6,18)
	resource_size_t p_last = -1;
#else
	__PTRDIFF_TYPE__ p_last = -1;
#endif

	DBG("Initilizing Dump...");

	if((err = setup())) {
		DBG("Setup Error");
		cleanup();
		return err;
	}

	
	for (p = iomem_resource.child; p ; p = p->sibling) {
		if (strncmp(p->name, LIME_RAMSTR, sizeof(LIME_RAMSTR)))
			continue;

		if (mode == LIME_MODE_LIME && (err = write_lime_header(p))) {
			DBG("Error writing header 0x%lx - 0x%lx", (long) p->start, (long) p->end);
			break;
		} else if (mode == LIME_MODE_PADDED && (err = write_padding((size_t) ((p->start - 1) - p_last)))) {
			DBG("Error writing padding 0x%lx - 0x%lx", (long) p_last, (long) p->start - 1);
			break;
		}

		if ((err = write_range(p))) {
			DBG("Error writing range 0x%lx - 0x%lx", (long) p->start, (long) p->end);
			break;
		}

		p_last = p->end;
	}
	

	cleanup();

	return err;
}

static int write_lime_header(struct resource * res) {
	long s;

	lime_mem_range_header header;

	memset(&header, 0, sizeof(lime_mem_range_header));
	header.magic = LIME_MAGIC;
	header.version = 1;
	header.s_addr = res->start;
	header.e_addr = res->end;
	
	s = write_vaddr(&header, sizeof(lime_mem_range_header));
	
	if (s != sizeof(lime_mem_range_header)) {
		DBG("Error sending header %ld", s);
		return (int) s;
	}				

	return 0;
}

static int write_padding(size_t s) {
	size_t i = 0;
	int r;

	while(s -= i) {

		i = min((size_t) PAGE_SIZE, s);
		r = write_vaddr(zero_page, i);

		if (r != i) {
			DBG("Error sending zero page: %d", r);
			return r;
		}
	}

	return 0;
}


static int write_range(struct resource * res) {
#if LINUX_VERSION_CODE >= KERNEL_VERSION(2,6,18)
	resource_size_t i, is, ds;
#else
	__PTRDIFF_TYPE__ i, is, ds;
#endif

    // for some reason, a page in the lowest 4GB of DRAM is missing a portion...
    // i.e. there is a single page that linux never fully writes out -- why??
    // fill the remainder of that one page with random data
    int  page_not_patched = 1;
    char randchars[16] = { 't', 'r', 'h', 'a',
                           'e', 'n', 'd', 'd',
                           'o', 'a', 'm', 'a',
                           'l', 'f', 'i', 'l'};
	struct page * p;
	void * v;
	
	int s;
    int krc;
    int urc;

	for (i = res->start; i <= res->end; i += PAGE_SIZE) {

		p = pfn_to_page((i) >> PAGE_SHIFT);
        krc = atomic_read( &p->_count );
        urc = atomic_read( &p->_mapcount );

        is = min((size_t) PAGE_SIZE, (size_t) (res->end - i + 1));
        ds = ( (size_t) PAGE_SIZE ) - is;

		v = kmap(p);
		s = write_vaddr(v, is);
		kunmap(p);
		
        if( page_not_patched ) {
            
            if( ds ) {
                // ds = page_size - bytes_in_memory_region
                while( ds > 16 ) {
                    s += write_vaddr( &randchars, 16 );
                    is += 16;
                    ds -= 16;
                }
                while( ds ) {
                    s += write_vaddr( &randchars, 1 );
                    is++;
                    ds--;
                }
                
                // this page patch is allowed to happen exactly once
                page_not_patched = 0;
            }
        }

		if( s != is ) {
			DBG("Error sending page %d", s);
			return (int) s;
		}
		
        s = write_count( krc, urc );
        if( s != sizeof( int ) ) {
            DBG("Error sending rc on page" );
            return (int) s;
        }
        pages_count++;
        if( pages_count == LIME_MAX_PAGES_PER_FILE ) {
            pages_count = 0;
            part_no++;
            cleanup_disk();
            setup_disk();
        }
	}

	return 0;
}

static int write_vaddr(void * v, size_t is) {
    return write_vaddr_disk(v, is);
}

static int setup(void) {
    return setup_disk();
}

static void cleanup(void) {
    return cleanup_disk();
}

void cleanup_module(void)
{
	
}

//MODULE_AUTHOR ("Joe T. Sylve, joe@digitalforensicssolutions.com");
//MODULE_DESCRIPTION ("Perform physical memory dump on Linux and Android devices.");
MODULE_LICENSE("GPL");
